from ultralytics import YOLO

# Load a model
model = YOLO("yolov8s.yaml")                            # 从头训练
# model = YOLO("pre_models/yolov8s.pt")                 # 加载预训练模型训练    
# model = YOLO('runs/detect/train/weights/best.pt')     # 从断点训练，resume设为True，默认为False

# Use the model
# model.train(data="./ultralytics/datasets/iruav.yaml", batch=128, lr0=0.04, epochs=100, device=[2,3], patience=20, save_json=True, resume=False, cube=True, dwa=False, val=True, gray=True, mod=False)
model.train(data="./ultralytics/datasets/antiuav310.yaml", batch=128, lr0=0.04, epochs=50, device=[0,1], patience=20, save_json=False, resume=False, cube=True, dwa=False, val=True, gray=True, mod=False)
# model.train(data="./ultralytics/datasets/uavswarm.yaml", batch=128, lr0=0.02, epochs=300, device=[0,1], patience=20, save_json=True, resume=False, cube=True, dwa=False, val=True, gray=True, mod=False)